from contextlib import AsyncExitStack
from typing import Any, Optional
from mcp.client.stdio import stdio_client
from mcp import ClientSession, StdioServerParameters, Tool, types
import asyncio

server_params = StdioServerParameters(
    command="python",
    args=["./mcp_server.py"],
    env=None
)


class MCPClient:
    def __init__(
        self,
        name: str,
        command: str,
        args: list[str],
        version: str = "0.01"
    ) -> None:
        self.session: Optional[ClientSession] = None
        self.exit_stack = AsyncExitStack()
        self.name = name
        self.command = command
        self.args = args
        self.tools: list[Tool] = []

    async def init(self) -> None:
        await self._connect_to_server()

    async def cleanup(self) -> None:
        try:
            await self.exit_stack.aclose()
        except Exception:
            print("Error during MCP client cleanup, traceback and continue!")
            # RICH_CONSOLE.print_exception()

    def get_tools(self) -> list[Tool]:
        return self.tools

    async def _connect_to_server(
        self,
    ) -> None:
        """
        Connect to an MCP server
        """
        server_params = StdioServerParameters(
            command=self.command,
            args=self.args,
        )

        stdio_transport = await self.exit_stack.enter_async_context(
            stdio_client(server_params),
        )
        self.stdio, self.write = stdio_transport
        self.session = await self.exit_stack.enter_async_context(
            ClientSession(self.stdio, self.write)
        )

        await self.session.initialize()

        # List available tools
        response = await self.session.list_tools()
        self.tools = response.tools
        print("\nConnected to server with tools:", [tool.name for tool in self.tools])

    async def call_tool(self, name: str, params: dict[str, Any]):
        return await self.session.call_tool(name, params)

async def example() -> None:
    client = MCPClient(
        name="calculate",
        command="python",
        args=["./mcp_server.py"],
    )
    await client.init()
    tools = client.get_tools()
    for tool in tools:
        print(tool.name)
    await client.cleanup()

if __name__ == "__main__":
    asyncio.run(example())